import numpy as np
import pandas as pd
import time
from span import Span

from aliyun.log import *


def log(msg):
    cur_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
    print(f'{cur_time}: {msg}')


def get_time_range_by_frag(frag):
    if frag == 1:
        return 1684800000, 1687392000  # 2023-05-23 08:00:00 ~ 2023-06-22 08:00:00
    elif frag == 2:
        return 1689984000, 1692576000  # 2023-07-22 08:00:00 ~ 2023-08-21 08:00:00
    elif frag == 3:
        return 1692576000, 1695254400  # 2023-08-21 08:00:00 ~ 2023-09-21 08:00:00
    else:
        return None, None


def exec_sls_query(query, db, start=1684800000, end=1695254400):
    project_name = "competition-1890508550122136"
    logstore_name = "unsatisfied_traces" if db == "unsatisfied_traces" else f"competition_2023-{db}"
    endpoint = "cn-guangzhou.log.aliyuncs.com"
    ak_id = "LTAI5tFq9GkTrh5gEZqjeijs"
    ak_key = "a3Ftbs5ges7W2tt5ohMLkSDe5SD7fe"
    sls_client = LogClient(endpoint, ak_id, ak_key)
    query_list = sls_client.get_log(project_name, logstore_name, start, end, "", query)

    results = []
    for result in query_list.get_logs():
        results.append(result.get_contents())

    return pd.DataFrame(results)


def fetch_trace_by_id(traceID):
    query = f"* | select traceID, spanID, parentSpanID, service, name, start, \"end\", duration from log where traceID='{traceID}' limit 1000000"
    trace = exec_sls_query(query, 'traces').astype({'start': np.int64, 'end': np.int64, 'duration': np.int64})

    print(f'#spans: {len(trace)}')
    assert check_trace(trace), "trace不完整"

    return build_trace(trace)


def check_trace(trace):
    # 检查是否有根节点
    if len(trace[trace['parentSpanID'] == '']) == 0:
        return False
    for idx, row in trace.iterrows():
        # 检查非根节点的父亲是否能找到
        if row['parentSpanID'] != '' and len(trace[trace['spanID'] == row['parentSpanID']]) == 0:
            return False
    return True


def build_trace(trace):
    trace = trace.drop_duplicates()
    rootSpanDf = trace[trace['parentSpanID'] == ''].iloc[0]
    rootSpan = Span(rootSpanDf['spanID'], rootSpanDf['service'], rootSpanDf['name'],
                    rootSpanDf['start'], rootSpanDf['end'], rootSpanDf['duration'])

    spanMap = {}
    curSpans = [rootSpan]
    while len(curSpans) > 0:
        nextCurSpans = []
        # 遍历该层所有 span
        for curSpan in curSpans:
            spanMap[curSpan.spanID] = curSpan
            # 找出当前 span 的所有子 span
            childrenSpanDf = trace[trace['parentSpanID'] == curSpan.spanID]
            # 将子 span 按照 start 排序
            childrenSpanDf = childrenSpanDf.sort_values('start')
            sum_children_duration = 0
            last_span_start, last_span_end = 0, 0
            # 遍历当前 span 的所有子 span
            for idx, spanData in childrenSpanDf.iterrows():
                childSpan = Span(spanData['spanID'], spanData['service'], spanData['name'],
                                 spanData['start'], spanData['end'], spanData['duration'])
                # 构建 span 父子关系
                childSpan.parent = curSpan
                curSpan.children.append(childSpan)
                # 计算 level
                childSpan.level = curSpan.level + 1
                # 计算 sum_children_duration（要考虑子 span 重合的情况）
                if childSpan.start < last_span_end:
                    last_span_end = max(last_span_end, childSpan.end)
                else:
                    sum_children_duration += last_span_end - last_span_start
                    last_span_start, last_span_end = childSpan.start, childSpan.end
                # 下一层的待展开 span
                nextCurSpans.append(childSpan)
            sum_children_duration += last_span_end - last_span_start
            # 计算 self_duration
            curSpan.self_duration = curSpan.duration - sum_children_duration

        curSpans = nextCurSpans

    return spanMap


def spanMap2Df(traceID, spanMap):
    return pd.DataFrame([[traceID, span.spanID, span.parent.spanID if span.parent else '', span.service, span.name,
                          span.start, span.end, span.duration, span.self_duration, span.level]
                         for span in list(spanMap.values())],

                        columns=['traceID', 'spanID', 'parentSpanID', 'service', 'name',
                                 'start', 'end', 'duration', 'self_duration', 'level'])